package nachos.threads;

import nachos.machine.*;

import java.util.Set;
import java.util.HashSet;
import java.util.Iterator;
import java.util.Random;
import java.util.Formatter;

/**
 * A scheduler that chooses threads using a lottery.
 *
 * <p>
 * A lottery scheduler associates a number of tickets with each thread. When a
 * thread needs to be dequeued, a random lottery is held, among all the tickets
 * of all the threads waiting to be dequeued. The thread that holds the winning
 * ticket is chosen.
 *
 * <p>
 * Note that a lottery scheduler must be able to handle a lot of tickets
 * (sometimes billions), so it is not acceptable to maintain state for every
 * ticket.
 *
 * <p>
 * A lottery scheduler must partially solve the priority inversion problem; in
 * particular, tickets must be transferred through locks, and through joins.
 * Unlike a priority scheduler, these tickets add (as opposed to just taking
 * the maximum).
 */
public class LotteryScheduler extends BasePriorityScheduler {

    /**
     * Allocate a new lottery scheduler.
     */
    public LotteryScheduler() {
	super(1, 1, Integer.MAX_VALUE);
    }

    /**
     * Allocate a new lottery thread queue.
     *
     * @param	transferPriority	<tt>true</tt> if this queue should
     *					transfer tickets from waiting threads
     *					to the owning thread.
     * @return	a new lottery thread queue.
     */
    public ThreadQueue newThreadQueue(boolean transferPriority) {
	return new LotteryQueue(transferPriority);
    }

    public ThreadState newThreadState(KThread thread) {
	return new ThreadState(thread);
    }

    /**
     * Maintains a set of ThreadStates that are waiting for a
     * resource, and the current owner of the resource.
     */
    private class LotteryQueue extends BasePriorityScheduler.PriorityQueue<ThreadState>
    {
	
	private Random rand = new Random();

	public Iterator<ThreadState> iterator() {
	    return threads.iterator();
	}

	public LotteryQueue(boolean transferPriority) {
	    super(transferPriority);
	}

	/**
	 * Return the total number of tickets (effective priority)
	 * held by all waiters.
	 */
	public int getTotalTickets() {
	    int tickets = 0;
	    for (ThreadState thread : this) {
		tickets += thread.getEffectivePriority();
	    }
	    return tickets;
	}

	public String toString() {
	    Formatter f = new Formatter();
	    f.format(super.toString());
	    f.format("  Total tickets %3d: ", getTotalTickets());
	    return f.toString();
	}

	public void addThread(ThreadState t) {
	    threads.add(t);
	}

	/**
	 * Remove a thread from the queue based on a lottery
	 * drawing. For each ThreadState, we know the number of
	 * tickets that it has. We choose a random number between 0
	 * (inclusive) and the total number of tickets held by any of
	 * my waiting threads. Iterate over the tickets and build a
	 * cumulative sum of the number of tickets seen so far, and
	 * stop when we've seen more tickets than the random number we
	 * chose. This will make it so that the probability of
	 * choosing a given thread is the same as the number of
	 * tickets held by that thread divided by the total number of
	 * tickets.
	 */
	@Override
	public ThreadState pop() {

	    ThreadState winner = null;

	    // Pick the winning ticket
	    int winningTicket = rand.nextInt(getTotalTickets());
	    
	    // Find the winner and remove it from my set of threads
	    int ticketsSoFar = 0;
	    
	    for (ThreadState t : this) {
		ticketsSoFar += t.getEffectivePriority();
		if (ticketsSoFar > winningTicket) {
		    winner = t;
		    break;
		}
	    }

	    Lib.assertTrue(winner != null);
	    threads.remove(winner);
	    return winner;
	}

	private Set<ThreadState> threads = new HashSet<ThreadState>();
    }

    protected class ThreadState 
	extends BasePriorityScheduler.ThreadState<LotteryQueue> {

	private LotteryQueue waitQueue;

	public ThreadState(KThread thread) {
	    super(thread);
	}

	/**
	 * To calculate the effective priority, just add up my
	 * priority and the total number of tickets in any queue that
	 * I own (that has transferPriority set to true).
	 */
	public void calculateEffectivePriority() {
	    effectivePriority = priority;
	    for (LotteryQueue q : ownedQueues) {
		if (q.transferPriority) {
		    effectivePriority += q.getTotalTickets();
		}
	    }
	}

	@Override
	public void setPriority(int priority) {
	    this.priority = priority;
	    calculateEffectivePriority();
	}

    }

    /**
     * Spawn off a few threads with different priority. Keep a global
     * counter of the number of work units remaining. Each of the
     * child threads will repeatedly decrement the work unit counter
     * and increment its own workDone counter. At the end, the threads
     * with higher priority should be able to do more work than
     * threads with lower priority.
     */
    public static void selfTest() {
	System.out.println("Testing lottery scheduler");
	final Lock lock = new Lock();
	
	class Work {
	    int units = 1000;
	}

	final Work work = new Work();

	class Worker implements Runnable {
	    int workDone = 0;
	    public void run() {
		boolean running = true;
		while (running) {
		    lock.acquire();
		    if (work.units > 0) {
			workDone++;
			work.units--;
		    }
		    else {
			running = false;
		    }
		    lock.release();
		    KThread.currentThread().yield();
		}
	    }
	}

	Worker[] workers = new Worker[5];
	KThread[] threads = new KThread[5];

	for (int i = 0; i < workers.length; i++) {
	    workers[i] = new Worker();
	    threads[i] = new KThread(workers[i]);
	    threads[i].setName(String.valueOf(i));
	}
	
	boolean intStatus = Machine.interrupt().disable();
	ThreadedKernel.scheduler.setPriority(threads[0], 1);
	ThreadedKernel.scheduler.setPriority(threads[1], 2);
	ThreadedKernel.scheduler.setPriority(threads[2], 4);
	ThreadedKernel.scheduler.setPriority(threads[3], 8);
	ThreadedKernel.scheduler.setPriority(threads[4], 16);

	Machine.interrupt().restore(intStatus);

	for (int i = 0; i < workers.length; i++) {
	    threads[i].fork();
	}

	for (int i = 0; i < workers.length; i++) {
	    threads[i].join();
	    System.out.printf("Worker %d did %5d%n", i, workers[i].workDone);
	    if (i > 0) {
		Lib.assertTrue(workers[i].workDone > workers[i-1].workDone);
	    }
	}

    }

}
